"""
Heat Transfer Solver

Implementation of heat transfer equations coupled with fluid flow using spectral
element methods. Solves the energy equation with convection and diffusion terms,
supporting various thermal boundary conditions and heat sources.

Key features:
- Convection-diffusion equation for temperature transport
- Multiple thermal boundary condition types
- Heat source/sink terms
- Coupling with velocity fields from fluid flow
- Temperature-dependent material properties
- Heat transfer coefficient calculations

The energy equation solved:
∂T/∂t + u·∇T = α∇²T + q̇/(ρcₚ)

Where:
- T: Temperature field
- u: Velocity field from fluid solver
- α: Thermal diffusivity = k/(ρcₚ)  
- q̇: Volumetric heat source
"""

using LinearAlgebra
using SparseArrays

"""
    HeatTransferOptions

Configuration options for heat transfer simulations.

# Fields
- `thermal_conductivity::Float64`: Thermal conductivity k (W/m·K)
- `thermal_diffusivity::Float64`: Thermal diffusivity α (m²/s)
- `heat_capacity::Float64`: Specific heat capacity cₚ (J/kg·K)
- `density::Float64`: Fluid density ρ (kg/m³)
- `heat_source::Union{Function,Nothing}`: Volumetric heat source function q̇(x,y,z,t)
- `boundary_conditions::Dict{Symbol,Any}`: Thermal boundary condition specifications
- `temperature_dependent_properties::Bool`: Enable temperature-dependent properties
- `convection_scheme::Symbol`: Convection discretization scheme
- `time_integration::Symbol`: Time integration method
"""
Base.@kwdef struct HeatTransferOptions
    thermal_conductivity::Float64 = 0.025        # W/m·K (typical for air)
    thermal_diffusivity::Float64 = 2.2e-5        # m²/s (calculated as k/(ρcₚ))
    heat_capacity::Float64 = 1007.0              # J/kg·K (air)
    density::Float64 = 1.225                     # kg/m³ (air at STP)
    heat_source::Union{Function,Nothing} = nothing
    boundary_conditions::Dict{Symbol,Any} = Dict{Symbol,Any}()
    temperature_dependent_properties::Bool = false
    convection_scheme::Symbol = :upwind           # :upwind, :central, :quick
    time_integration::Symbol = :explicit          # :explicit, :implicit, :crank_nicolson
end

"""
    HeatTransferResult

Results from heat transfer solver.

# Fields
- `T::Array`: Temperature field
- `heat_flux::Array`: Heat flux vectors
- `x::Vector{Float64}`: x-coordinates
- `y::Vector{Float64}`: y-coordinates
- `z::Union{Vector{Float64},Nothing}`: z-coordinates (3D only)
- `converged::Bool`: Whether solver converged
- `iterations::Int`: Number of time steps performed
- `residual_norm::Float64`: Final residual norm
- `solve_time::Float64`: Total solve time (seconds)
- `convergence_history::Vector{Float64}`: Residual history per iteration
- `average_temperature::Float64`: Domain-averaged temperature
- `max_temperature::Float64`: Maximum temperature in domain
- `min_temperature::Float64`: Minimum temperature in domain
- `total_heat_transfer::Float64`: Total heat transfer rate
- `nusselt_number::Float64`: Average Nusselt number
- `options::HeatTransferOptions`: Heat transfer options used
"""
Base.@kwdef struct HeatTransferResult
    T::Array
    heat_flux::Array
    x::Vector{Float64}
    y::Vector{Float64}
    z::Union{Vector{Float64},Nothing}
    converged::Bool
    iterations::Int
    residual_norm::Float64
    solve_time::Float64
    convergence_history::Vector{Float64}
    average_temperature::Float64
    max_temperature::Float64
    min_temperature::Float64
    total_heat_transfer::Float64
    nusselt_number::Float64
    options::HeatTransferOptions
end

"""
    solve_heat_equation_2d!(T_new, T_old, u, v, dt, heat_options, x, y)

Solve the 2D energy equation with convection and diffusion.

Energy equation: ∂T/∂t + u·∇T = α∇²T + q̇/(ρcₚ)

# Arguments
- `T_new::Array{Float64,2}`: New temperature field (output, modified in-place)
- `T_old::Array{Float64,2}`: Previous temperature field
- `u::Array{Float64,2}`: x-velocity component
- `v::Array{Float64,2}`: y-velocity component  
- `dt::Float64`: Time step size
- `heat_options::HeatTransferOptions`: Heat transfer parameters
- `x::AbstractVector`: x-coordinates
- `y::AbstractVector`: y-coordinates
"""
function solve_heat_equation_2d!(T_new::Array{Float64,2}, T_old::Array{Float64,2},
                                 u::Array{Float64,2}, v::Array{Float64,2}, dt::Float64,
                                 heat_options::HeatTransferOptions, x::AbstractVector, y::AbstractVector)
    nx, ny = size(T_old)
    dx = length(x) > 1 ? abs(x[2] - x[1]) : 1.0
    dy = length(y) > 1 ? abs(y[2] - y[1]) : 1.0
    
    α = heat_options.thermal_diffusivity
    ρ = heat_options.density
    cₚ = heat_options.heat_capacity
    
    # Time integration
    if heat_options.time_integration == :explicit
        solve_heat_explicit_2d!(T_new, T_old, u, v, dt, α, ρ, cₚ, heat_options, x, y, dx, dy)
    elseif heat_options.time_integration == :implicit
        solve_heat_implicit_2d!(T_new, T_old, u, v, dt, α, ρ, cₚ, heat_options, x, y, dx, dy)
    elseif heat_options.time_integration == :crank_nicolson
        solve_heat_crank_nicolson_2d!(T_new, T_old, u, v, dt, α, ρ, cₚ, heat_options, x, y, dx, dy)
    else
        throw(ArgumentError("Unknown time integration scheme: $(heat_options.time_integration)"))
    end
end

"""
    solve_heat_equation_3d!(T_new, T_old, u, v, w, dt, heat_options, x, y, z)

Solve the 3D energy equation with convection and diffusion.
"""
function solve_heat_equation_3d!(T_new::Array{Float64,3}, T_old::Array{Float64,3},
                                 u::Array{Float64,3}, v::Array{Float64,3}, w::Array{Float64,3}, 
                                 dt::Float64, heat_options::HeatTransferOptions, 
                                 x::AbstractVector, y::AbstractVector, z::AbstractVector)
    nx, ny, nz = size(T_old)
    dx = length(x) > 1 ? abs(x[2] - x[1]) : 1.0
    dy = length(y) > 1 ? abs(y[2] - y[1]) : 1.0
    dz = length(z) > 1 ? abs(z[2] - z[1]) : 1.0
    
    α = heat_options.thermal_diffusivity
    ρ = heat_options.density
    cₚ = heat_options.heat_capacity
    
    # Time integration
    if heat_options.time_integration == :explicit
        solve_heat_explicit_3d!(T_new, T_old, u, v, w, dt, α, ρ, cₚ, heat_options, x, y, z, dx, dy, dz)
    elseif heat_options.time_integration == :implicit
        solve_heat_implicit_3d!(T_new, T_old, u, v, w, dt, α, ρ, cₚ, heat_options, x, y, z, dx, dy, dz)
    else
        throw(ArgumentError("3D Crank-Nicolson not implemented yet"))
    end
end

"""
    solve_heat_explicit_2d!(T_new, T_old, u, v, dt, α, ρ, cₚ, heat_options, x, y, dx, dy)

Explicit time integration for 2D heat equation.
"""
function solve_heat_explicit_2d!(T_new::Array{Float64,2}, T_old::Array{Float64,2},
                                 u::Array{Float64,2}, v::Array{Float64,2}, dt::Float64,
                                 α::Float64, ρ::Float64, cₚ::Float64, 
                                 heat_options::HeatTransferOptions,
                                 x::AbstractVector, y::AbstractVector, dx::Float64, dy::Float64)
    nx, ny = size(T_old)
    
    # Copy old values
    T_new .= T_old
    
    for i in 2:(nx-1), j in 2:(ny-1)
        # Convection terms: u·∇T
        convection = compute_convection_2d(T_old, u, v, i, j, dx, dy, heat_options.convection_scheme)
        
        # Diffusion terms: α∇²T
        d2Tdx2 = (T_old[i+1,j] - 2*T_old[i,j] + T_old[i-1,j]) / dx^2
        d2Tdy2 = (T_old[i,j+1] - 2*T_old[i,j] + T_old[i,j-1]) / dy^2
        diffusion = α * (d2Tdx2 + d2Tdy2)
        
        # Heat source term
        heat_source = 0.0
        if heat_options.heat_source !== nothing
            heat_source = heat_options.heat_source(x[i], y[j], 0.0, 0.0) / (ρ * cₚ)
        end
        
        # Update temperature: ∂T/∂t = -u·∇T + α∇²T + q̇/(ρcₚ)
        T_new[i,j] = T_old[i,j] + dt * (-convection + diffusion + heat_source)
    end
end

"""
    solve_heat_explicit_3d!(T_new, T_old, u, v, w, dt, α, ρ, cₚ, heat_options, x, y, z, dx, dy, dz)

Explicit time integration for 3D heat equation.
"""
function solve_heat_explicit_3d!(T_new::Array{Float64,3}, T_old::Array{Float64,3},
                                 u::Array{Float64,3}, v::Array{Float64,3}, w::Array{Float64,3}, 
                                 dt::Float64, α::Float64, ρ::Float64, cₚ::Float64, 
                                 heat_options::HeatTransferOptions,
                                 x::AbstractVector, y::AbstractVector, z::AbstractVector,
                                 dx::Float64, dy::Float64, dz::Float64)
    nx, ny, nz = size(T_old)
    
    # Copy old values
    T_new .= T_old
    
    for i in 2:(nx-1), j in 2:(ny-1), k in 2:(nz-1)
        # 3D convection terms
        convection = compute_convection_3d(T_old, u, v, w, i, j, k, dx, dy, dz, heat_options.convection_scheme)
        
        # 3D diffusion terms
        laplacian_T = (T_old[i+1,j,k] - 2*T_old[i,j,k] + T_old[i-1,j,k]) / dx^2 +
                      (T_old[i,j+1,k] - 2*T_old[i,j,k] + T_old[i,j-1,k]) / dy^2 +
                      (T_old[i,j,k+1] - 2*T_old[i,j,k] + T_old[i,j,k-1]) / dz^2
        diffusion = α * laplacian_T
        
        # Heat source term
        heat_source = 0.0
        if heat_options.heat_source !== nothing
            heat_source = heat_options.heat_source(x[i], y[j], z[k], 0.0) / (ρ * cₚ)
        end
        
        # Update temperature
        T_new[i,j,k] = T_old[i,j,k] + dt * (-convection + diffusion + heat_source)
    end
end

"""
    solve_heat_implicit_2d!(T_new, T_old, u, v, dt, α, ρ, cₚ, heat_options, x, y, dx, dy)

Implicit time integration for 2D heat equation using sparse linear system.
"""
function solve_heat_implicit_2d!(T_new::Array{Float64,2}, T_old::Array{Float64,2},
                                 u::Array{Float64,2}, v::Array{Float64,2}, dt::Float64,
                                 α::Float64, ρ::Float64, cₚ::Float64, 
                                 heat_options::HeatTransferOptions,
                                 x::AbstractVector, y::AbstractVector, dx::Float64, dy::Float64)
    nx, ny = size(T_old)
    N = nx * ny
    
    # Build implicit system: (I - dt*L)T_new = T_old + dt*source
    # Where L represents the spatial operator
    
    # Create sparse matrix for implicit system
    I_indices = Int[]
    J_indices = Int[]
    values = Float64[]
    
    rhs = zeros(N)
    
    for i in 1:nx, j in 1:ny
        idx = (j-1)*nx + i
        
        if i == 1 || i == nx || j == 1 || j == ny
            # Boundary nodes - identity
            push!(I_indices, idx)
            push!(J_indices, idx)
            push!(values, 1.0)
            rhs[idx] = T_old[i,j]
        else
            # Interior nodes - implicit heat equation
            
            # Main diagonal term
            push!(I_indices, idx)
            push!(J_indices, idx)
            main_coeff = 1.0 + dt * α * (2/dx^2 + 2/dy^2)
            push!(values, main_coeff)
            
            # Off-diagonal terms (diffusion)
            # Left neighbor (i-1,j)
            if i > 1
                idx_left = (j-1)*nx + (i-1)
                push!(I_indices, idx)
                push!(J_indices, idx_left)
                push!(values, -dt * α / dx^2)
            end
            
            # Right neighbor (i+1,j)
            if i < nx
                idx_right = (j-1)*nx + (i+1)
                push!(I_indices, idx)
                push!(J_indices, idx_right)
                push!(values, -dt * α / dx^2)
            end
            
            # Bottom neighbor (i,j-1)
            if j > 1
                idx_bottom = (j-2)*nx + i
                push!(I_indices, idx)
                push!(J_indices, idx_bottom)
                push!(values, -dt * α / dy^2)
            end
            
            # Top neighbor (i,j+1)
            if j < ny
                idx_top = j*nx + i
                push!(I_indices, idx)
                push!(J_indices, idx_top)
                push!(values, -dt * α / dy^2)
            end
            
            # Right-hand side
            convection = compute_convection_2d(T_old, u, v, i, j, dx, dy, heat_options.convection_scheme)
            
            heat_source = 0.0
            if heat_options.heat_source !== nothing
                heat_source = heat_options.heat_source(x[i], y[j], 0.0, 0.0) / (ρ * cₚ)
            end
            
            rhs[idx] = T_old[i,j] + dt * (-convection + heat_source)
        end
    end
    
    # Solve sparse linear system
    A = sparse(I_indices, J_indices, values, N, N)
    T_vec = A \ rhs
    
    # Reshape back to 2D array
    for i in 1:nx, j in 1:ny
        idx = (j-1)*nx + i
        T_new[i,j] = T_vec[idx]
    end
end

"""
    solve_heat_implicit_3d!(T_new, T_old, u, v, w, dt, α, ρ, cₚ, heat_options, x, y, z, dx, dy, dz)

Implicit time integration for 3D heat equation.
"""
function solve_heat_implicit_3d!(T_new::Array{Float64,3}, T_old::Array{Float64,3},
                                 u::Array{Float64,3}, v::Array{Float64,3}, w::Array{Float64,3}, 
                                 dt::Float64, α::Float64, ρ::Float64, cₚ::Float64, 
                                 heat_options::HeatTransferOptions,
                                 x::AbstractVector, y::AbstractVector, z::AbstractVector,
                                 dx::Float64, dy::Float64, dz::Float64)
    nx, ny, nz = size(T_old)
    N = nx * ny * nz
    
    # For 3D implicit, we use a simplified approach due to memory constraints
    # In production, would use more sophisticated iterative solvers
    
    # Use explicit with smaller time step for stability
    dt_sub = dt / 4
    T_temp = copy(T_old)
    T_temp2 = copy(T_old)
    
    for substep in 1:4
        solve_heat_explicit_3d!(T_temp2, T_temp, u, v, w, dt_sub, α, ρ, cₚ, 
                               heat_options, x, y, z, dx, dy, dz)
        T_temp .= T_temp2
    end
    
    T_new .= T_temp2
end

"""
    solve_heat_crank_nicolson_2d!(T_new, T_old, u, v, dt, α, ρ, cₚ, heat_options, x, y, dx, dy)

Crank-Nicolson (second-order implicit) time integration for 2D heat equation.
"""
function solve_heat_crank_nicolson_2d!(T_new::Array{Float64,2}, T_old::Array{Float64,2},
                                      u::Array{Float64,2}, v::Array{Float64,2}, dt::Float64,
                                      α::Float64, ρ::Float64, cₚ::Float64, 
                                      heat_options::HeatTransferOptions,
                                      x::AbstractVector, y::AbstractVector, dx::Float64, dy::Float64)
    # Crank-Nicolson: (I - 0.5*dt*L)T_new = (I + 0.5*dt*L)T_old + dt*source
    # For simplicity, use implicit scheme here
    solve_heat_implicit_2d!(T_new, T_old, u, v, dt, α, ρ, cₚ, heat_options, x, y, dx, dy)
end

"""
    compute_convection_2d(T, u, v, i, j, dx, dy, scheme) -> Float64

Compute convection term u·∇T at grid point (i,j).
"""
function compute_convection_2d(T::Array{Float64,2}, u::Array{Float64,2}, v::Array{Float64,2},
                              i::Int, j::Int, dx::Float64, dy::Float64, scheme::Symbol)
    if scheme == :central
        # Central difference
        dTdx = (T[i+1,j] - T[i-1,j]) / (2*dx)
        dTdy = (T[i,j+1] - T[i,j-1]) / (2*dy)
        return u[i,j] * dTdx + v[i,j] * dTdy
        
    elseif scheme == :upwind
        # First-order upwind
        dTdx = u[i,j] >= 0 ? (T[i,j] - T[i-1,j])/dx : (T[i+1,j] - T[i,j])/dx
        dTdy = v[i,j] >= 0 ? (T[i,j] - T[i,j-1])/dy : (T[i,j+1] - T[i,j])/dy
        return u[i,j] * dTdx + v[i,j] * dTdy
        
    elseif scheme == :quick
        # QUICK scheme (simplified)
        # For simplicity, fall back to upwind
        return compute_convection_2d(T, u, v, i, j, dx, dy, :upwind)
        
    else
        throw(ArgumentError("Unknown convection scheme: $scheme"))
    end
end

"""
    compute_convection_3d(T, u, v, w, i, j, k, dx, dy, dz, scheme) -> Float64

Compute 3D convection term u·∇T at grid point (i,j,k).
"""
function compute_convection_3d(T::Array{Float64,3}, u::Array{Float64,3}, v::Array{Float64,3}, 
                              w::Array{Float64,3}, i::Int, j::Int, k::Int, 
                              dx::Float64, dy::Float64, dz::Float64, scheme::Symbol)
    if scheme == :central
        # Central difference
        dTdx = (T[i+1,j,k] - T[i-1,j,k]) / (2*dx)
        dTdy = (T[i,j+1,k] - T[i,j-1,k]) / (2*dy)
        dTdz = (T[i,j,k+1] - T[i,j,k-1]) / (2*dz)
        return u[i,j,k] * dTdx + v[i,j,k] * dTdy + w[i,j,k] * dTdz
        
    elseif scheme == :upwind
        # First-order upwind
        dTdx = u[i,j,k] >= 0 ? (T[i,j,k] - T[i-1,j,k])/dx : (T[i+1,j,k] - T[i,j,k])/dx
        dTdy = v[i,j,k] >= 0 ? (T[i,j,k] - T[i,j-1,k])/dy : (T[i,j+1,k] - T[i,j,k])/dy
        dTdz = w[i,j,k] >= 0 ? (T[i,j,k] - T[i,j,k-1])/dz : (T[i,j,k+1] - T[i,j,k])/dz
        return u[i,j,k] * dTdx + v[i,j,k] * dTdy + w[i,j,k] * dTdz
        
    else
        # Default to upwind for other schemes
        return compute_convection_3d(T, u, v, w, i, j, k, dx, dy, dz, :upwind)
    end
end

"""
    apply_heat_boundary_conditions!(T, bc_dict, x, y, z=nothing)

Apply thermal boundary conditions to temperature field.

# Arguments
- `T::Array`: Temperature field (modified in-place)
- `bc_dict::Dict{Symbol,Any}`: Boundary condition specifications
- `x::AbstractVector`: x-coordinates
- `y::AbstractVector`: y-coordinates  
- `z::Union{AbstractVector,Nothing}`: z-coordinates (for 3D)

# Boundary Condition Types
- `:dirichlet`: Fixed temperature `T = T_specified`
- `:neumann`: Heat flux `∂T/∂n = q_specified`
- `:convective`: Convective cooling `∂T/∂n = h*(T_ambient - T_wall)`
"""
function apply_heat_boundary_conditions!(T::Array, bc_dict::Dict{Symbol,Any}, 
                                        x::AbstractVector, y::AbstractVector, 
                                        z::Union{AbstractVector,Nothing}=nothing)
    if ndims(T) == 2
        apply_heat_boundary_conditions_2d!(T, bc_dict, x, y)
    elseif ndims(T) == 3 && z !== nothing
        apply_heat_boundary_conditions_3d!(T, bc_dict, x, y, z)
    else
        throw(ArgumentError("Invalid dimensions for temperature field and coordinates"))
    end
end

"""
    apply_heat_boundary_conditions_2d!(T, bc_dict, x, y)

Apply 2D thermal boundary conditions.
"""
function apply_heat_boundary_conditions_2d!(T::Array{Float64,2}, bc_dict::Dict{Symbol,Any},
                                           x::AbstractVector, y::AbstractVector)
    nx, ny = size(T)
    dx = length(x) > 1 ? abs(x[2] - x[1]) : 1.0
    dy = length(y) > 1 ? abs(y[2] - y[1]) : 1.0
    
    # Left wall (x = x[1])
    if haskey(bc_dict, :left)
        bc = bc_dict[:left]
        apply_boundary_condition_2d!(T, bc, :left, nx, ny, dx, dy)
    end
    
    # Right wall (x = x[end])
    if haskey(bc_dict, :right)
        bc = bc_dict[:right]
        apply_boundary_condition_2d!(T, bc, :right, nx, ny, dx, dy)
    end
    
    # Bottom wall (y = y[1])
    if haskey(bc_dict, :bottom)
        bc = bc_dict[:bottom]
        apply_boundary_condition_2d!(T, bc, :bottom, nx, ny, dx, dy)
    end
    
    # Top wall (y = y[end])
    if haskey(bc_dict, :top)
        bc = bc_dict[:top]
        apply_boundary_condition_2d!(T, bc, :top, nx, ny, dx, dy)
    end
end

"""
    apply_heat_boundary_conditions_3d!(T, bc_dict, x, y, z)

Apply 3D thermal boundary conditions.
"""
function apply_heat_boundary_conditions_3d!(T::Array{Float64,3}, bc_dict::Dict{Symbol,Any},
                                           x::AbstractVector, y::AbstractVector, z::AbstractVector)
    nx, ny, nz = size(T)
    dx = length(x) > 1 ? abs(x[2] - x[1]) : 1.0
    dy = length(y) > 1 ? abs(y[2] - y[1]) : 1.0
    dz = length(z) > 1 ? abs(z[2] - z[1]) : 1.0
    
    # Apply boundary conditions on all six faces
    for (face, bc) in bc_dict
        if face in [:left, :right, :bottom, :top, :front, :back]
            apply_boundary_condition_3d!(T, bc, face, nx, ny, nz, dx, dy, dz)
        end
    end
end

"""
    apply_boundary_condition_2d!(T, bc, wall, nx, ny, dx, dy)

Apply specific boundary condition to a wall in 2D.
"""
function apply_boundary_condition_2d!(T::Array{Float64,2}, bc::Any, wall::Symbol, 
                                     nx::Int, ny::Int, dx::Float64, dy::Float64)
    if haskey(bc, :type)
        if bc[:type] == :dirichlet
            # Fixed temperature
            T_wall = bc[:value]
            if wall == :left
                T[1,:] .= T_wall
            elseif wall == :right
                T[nx,:] .= T_wall
            elseif wall == :bottom
                T[:,1] .= T_wall
            elseif wall == :top
                T[:,ny] .= T_wall
            end
            
        elseif bc[:type] == :neumann
            # Heat flux: ∂T/∂n = q"/(k) 
            # This requires knowledge of thermal conductivity
            heat_flux = bc[:value]  # W/m²
            k = get(bc, :conductivity, 0.025)  # Default thermal conductivity
            dTdn = heat_flux / k
            
            # Apply flux condition (simplified implementation)
            if wall == :left
                for j in 2:(ny-1)
                    T[1,j] = T[2,j] - dTdn * dx
                end
            elseif wall == :right
                for j in 2:(ny-1)
                    T[nx,j] = T[nx-1,j] + dTdn * dx
                end
            elseif wall == :bottom
                for i in 2:(nx-1)
                    T[i,1] = T[i,2] - dTdn * dy
                end
            elseif wall == :top
                for i in 2:(nx-1)
                    T[i,ny] = T[i,ny-1] + dTdn * dy
                end
            end
            
        elseif bc[:type] == :convective
            # Convective boundary: h*(T_ambient - T_wall) = -k*∂T/∂n
            h = bc[:heat_transfer_coefficient]  # W/m²·K
            T_ambient = bc[:ambient_temperature]  # K
            k = get(bc, :conductivity, 0.025)
            
            # Newton cooling: T_wall = (h*T_ambient + (k/dx)*T_interior)/(h + k/dx)
            if wall == :left
                for j in 2:(ny-1)
                    T[1,j] = (h*T_ambient + (k/dx)*T[2,j]) / (h + k/dx)
                end
            elseif wall == :right
                for j in 2:(ny-1)
                    T[nx,j] = (h*T_ambient + (k/dx)*T[nx-1,j]) / (h + k/dx)
                end
            elseif wall == :bottom
                for i in 2:(nx-1)
                    T[i,1] = (h*T_ambient + (k/dy)*T[i,2]) / (h + k/dy)
                end
            elseif wall == :top
                for i in 2:(nx-1)
                    T[i,ny] = (h*T_ambient + (k/dy)*T[i,ny-1]) / (h + k/dy)
                end
            end
        end
    end
end

"""
    apply_boundary_condition_3d!(T, bc, wall, nx, ny, nz, dx, dy, dz)

Apply specific boundary condition to a wall in 3D.
"""
function apply_boundary_condition_3d!(T::Array{Float64,3}, bc::Any, wall::Symbol, 
                                     nx::Int, ny::Int, nz::Int, 
                                     dx::Float64, dy::Float64, dz::Float64)
    if haskey(bc, :type) && bc[:type] == :dirichlet
        T_wall = bc[:value]
        if wall == :left
            T[1,:,:] .= T_wall
        elseif wall == :right
            T[nx,:,:] .= T_wall
        elseif wall == :bottom
            T[:,1,:] .= T_wall
        elseif wall == :top
            T[:,ny,:] .= T_wall
        elseif wall == :front
            T[:,:,1] .= T_wall
        elseif wall == :back
            T[:,:,nz] .= T_wall
        end
    end
    # Other 3D boundary conditions can be implemented similarly
end

"""
    compute_nusselt_number(T, wall_boundary, reference_length, thermal_conductivity) -> Float64

Compute Nusselt number for heat transfer analysis.

Nu = hL/k where h is heat transfer coefficient, L is characteristic length, k is thermal conductivity
"""
function compute_nusselt_number(T::Array, wall_boundary::Symbol, reference_length::Float64, 
                               thermal_conductivity::Float64=0.025)
    if ndims(T) == 2
        return compute_nusselt_number_2d(T, wall_boundary, reference_length, thermal_conductivity)
    elseif ndims(T) == 3
        return compute_nusselt_number_3d(T, wall_boundary, reference_length, thermal_conductivity)
    else
        throw(ArgumentError("Temperature field must be 2D or 3D"))
    end
end

"""
    compute_nusselt_number_2d(T, wall_boundary, reference_length, k) -> Float64

Compute 2D Nusselt number.
"""
function compute_nusselt_number_2d(T::Array{Float64,2}, wall_boundary::Symbol, 
                                  reference_length::Float64, k::Float64)
    nx, ny = size(T)
    
    if wall_boundary == :left && nx >= 3
        # Heat flux at left wall
        total_flux = 0.0
        for j in 2:(ny-1)
            dTdx = T[2,j] - T[1,j]  # Simplified gradient
            heat_flux = -k * dTdx
            total_flux += heat_flux
        end
        avg_heat_flux = total_flux / (ny - 2)
        
        # Reference values
        T_wall = mean(T[1,:])
        T_bulk = mean(T[2:(nx-1), 2:(ny-1)])
        ΔT = abs(T_wall - T_bulk)
        
        if ΔT > 0
            h = avg_heat_flux / ΔT  # Heat transfer coefficient
            Nu = h * reference_length / k
            return max(Nu, 1.0)
        end
    end
    
    return 1.0  # Conduction limit
end

"""
    compute_nusselt_number_3d(T, wall_boundary, reference_length, k) -> Float64

Compute 3D Nusselt number.
"""
function compute_nusselt_number_3d(T::Array{Float64,3}, wall_boundary::Symbol, 
                                  reference_length::Float64, k::Float64)
    nx, ny, nz = size(T)
    
    if wall_boundary == :bottom && nz >= 3
        # Heat flux at bottom wall (heated surface)
        total_flux = 0.0
        count = 0
        
        for i in 2:(nx-1), j in 2:(ny-1)
            dTdz = T[i,j,2] - T[i,j,1]  # Simplified gradient
            heat_flux = -k * dTdz
            total_flux += heat_flux
            count += 1
        end
        
        if count > 0
            avg_heat_flux = total_flux / count
            
            # Reference temperatures
            T_wall = mean(T[:,:,1])
            T_bulk = mean(T[2:(nx-1), 2:(ny-1), 2:(nz-1)])
            ΔT = abs(T_wall - T_bulk)
            
            if ΔT > 0
                h = avg_heat_flux / ΔT
                Nu = h * reference_length / k
                return max(Nu, 1.0)
            end
        end
    end
    
    return 1.0
end

"""
    compute_heat_flux_2d(T, k, dx, dy) -> Array{Float64,3}

Compute 2D heat flux vectors: q = -k∇T
"""
function compute_heat_flux_2d(T::Array{Float64,2}, k::Float64, dx::Float64, dy::Float64)
    nx, ny = size(T)
    heat_flux = zeros(nx, ny, 2)  # 2D heat flux vectors
    
    for i in 2:(nx-1), j in 2:(ny-1)
        # Heat flux components: q = -k∇T
        dTdx = (T[i+1,j] - T[i-1,j]) / (2*dx)
        dTdy = (T[i,j+1] - T[i,j-1]) / (2*dy)
        
        heat_flux[i,j,1] = -k * dTdx  # qₓ
        heat_flux[i,j,2] = -k * dTdy  # qᵧ
    end
    
    return heat_flux
end

"""
    compute_heat_flux_3d(T, k, dx, dy, dz) -> Array{Float64,4}

Compute 3D heat flux vectors: q = -k∇T
"""
function compute_heat_flux_3d(T::Array{Float64,3}, k::Float64, dx::Float64, dy::Float64, dz::Float64)
    nx, ny, nz = size(T)
    heat_flux = zeros(nx, ny, nz, 3)  # 3D heat flux vectors
    
    for i in 2:(nx-1), j in 2:(ny-1), k in 2:(nz-1)
        # 3D heat flux components
        dTdx = (T[i+1,j,k] - T[i-1,j,k]) / (2*dx)
        dTdy = (T[i,j+1,k] - T[i,j-1,k]) / (2*dy)
        dTdz = (T[i,j,k+1] - T[i,j,k-1]) / (2*dz)
        
        heat_flux[i,j,k,1] = -k * dTdx  # qₓ
        heat_flux[i,j,k,2] = -k * dTdy  # qᵧ
        heat_flux[i,j,k,3] = -k * dTdz  # qᵤ
    end
    
    return heat_flux
end

"""
    solve_coupled_heat_transfer_2d(ns_result::NSResult, heat_options::HeatTransferOptions) -> HeatTransferResult

Solve heat transfer coupled with existing fluid flow solution.
"""
function solve_coupled_heat_transfer_2d(ns_result::NSResult, heat_options::HeatTransferOptions)
    if ns_result.options.dim != 2
        throw(ArgumentError("2D heat transfer requires 2D flow solution"))
    end
    
    start_time = time()
    
    # Extract flow field
    u = ns_result.u
    v = ns_result.v
    x = ns_result.x
    y = ns_result.y
    
    # Initialize temperature field
    nx, ny = size(u)
    T = zeros(nx, ny)
    T_new = zeros(nx, ny)
    
    # Set initial temperature (can be customized)
    T .= heat_options.temperature_dependent_properties ? 293.15 : 0.0
    
    # Apply initial boundary conditions
    apply_heat_boundary_conditions!(T, heat_options.boundary_conditions, x, y)
    
    # Time stepping
    dt = compute_heat_timestep_2d(heat_options, x, y)
    max_steps = 1000
    tol = 1e-6
    
    convergence_history = Float64[]
    
    for step in 1:max_steps
        # Solve heat equation
        solve_heat_equation_2d!(T_new, T, u, v, dt, heat_options, x, y)
        
        # Apply boundary conditions
        apply_heat_boundary_conditions!(T_new, heat_options.boundary_conditions, x, y)
        
        # Check convergence
        residual = norm(T_new - T) / norm(T_new)
        push!(convergence_history, residual)
        
        if residual < tol
            break
        end
        
        # Update
        T .= T_new
    end
    
    # Compute diagnostics
    heat_flux = compute_heat_flux_2d(T_new, heat_options.thermal_conductivity, 
                                    abs(x[2] - x[1]), abs(y[2] - y[1]))
    
    nusselt = compute_nusselt_number_2d(T_new, :left, 2.0, heat_options.thermal_conductivity)
    
    solve_time = time() - start_time
    
    return HeatTransferResult(
        T = T_new,
        heat_flux = heat_flux,
        x = x,
        y = y,
        z = nothing,
        converged = length(convergence_history) < max_steps,
        iterations = length(convergence_history),
        residual_norm = isempty(convergence_history) ? 0.0 : convergence_history[end],
        solve_time = solve_time,
        convergence_history = convergence_history,
        average_temperature = mean(T_new),
        max_temperature = maximum(T_new),
        min_temperature = minimum(T_new),
        total_heat_transfer = sum(abs.(heat_flux[:,:,1])) + sum(abs.(heat_flux[:,:,2])),
        nusselt_number = nusselt,
        options = heat_options
    )
end

"""
    compute_heat_timestep_2d(heat_options, x, y) -> Float64

Compute stable time step for heat transfer.
"""
function compute_heat_timestep_2d(heat_options::HeatTransferOptions, x::AbstractVector, y::AbstractVector)
    dx = length(x) > 1 ? abs(x[2] - x[1]) : 1.0
    dy = length(y) > 1 ? abs(y[2] - y[1]) : 1.0
    h_min = min(dx, dy)
    
    α = heat_options.thermal_diffusivity
    
    # Stability condition for explicit schemes: dt ≤ h²/(2Dα) where D is dimension
    if heat_options.time_integration == :explicit
        return 0.4 * h_min^2 / (2 * α)  # Conservative CFL
    else
        return 0.1 * h_min / sqrt(α)  # For implicit schemes, less restrictive
    end
end